#include "compi/compi.h"

#include <cstring>
#include <iostream>
#include <string>
#include <vector>

using namespace compi;

static std::vector<std::byte>
to_bytes(const std::string & s)
{
    std::vector<std::byte> v(s.size());
    std::memcpy(v.data(), s.data(), s.size());
    return v;
}

static std::string
to_string(const std::vector<std::byte> & v)
{
    return std::string(reinterpret_cast<const char *>(v.data()), v.size());
}

static void
demo_channel_explicit_tag(Channel & ch)
{
    int rank = ch.rank();
    int peer = (rank == 0 ? 1 : 0);
    const int tag_send = 42;
    const int tag_reply = 43;

    if (rank == 0)
    {
        std::vector<std::byte> payload = to_bytes("hello-channel");
        auto sres = ch.send_frame(peer, payload, tag_send);
        if (sres.is_error())
        {
            std::cerr << "[channel-explicit] send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-explicit] rank 0 sent: " << to_string(payload) << std::endl;

        std::vector<std::byte> reply;
        auto rres = ch.recv_frame(peer, reply, tag_reply);
        if (rres.is_error())
        {
            std::cerr << "[channel-explicit] recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-explicit] rank 0 received: " << to_string(reply) << std::endl;
    }
    else
    {
        std::vector<std::byte> received;
        auto rres = ch.recv_frame(peer, received, tag_send);
        if (rres.is_error())
        {
            std::cerr << "[channel-explicit] recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-explicit] rank 1 received: " << to_string(received) << std::endl;

        std::vector<std::byte> response = to_bytes(to_string(received) + " | ack");
        auto sres = ch.send_frame(peer, response, tag_reply);
        if (sres.is_error())
        {
            std::cerr << "[channel-explicit] send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-explicit] rank 1 sent: " << to_string(response) << std::endl;
    }
}

static void
demo_channel_namespace_any(Channel & ch, TagNamespace ns)
{
    int rank = ch.rank();
    int peer = (rank == 0 ? 1 : 0);

    if (rank == 0)
    {
        std::vector<std::byte> payload = to_bytes("ns-any: hi");
        auto sres = ch.send_frame(peer, payload, ns);
        if (sres.is_error())
        {
            std::cerr << "[channel-ns] send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-ns] rank 0 sent: " << to_string(payload) << std::endl;

        std::vector<std::byte> reply;
        auto rres = ch.recv_frame(peer, reply, ns);
        if (rres.is_error())
        {
            std::cerr << "[channel-ns] recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-ns] rank 0 received: " << to_string(reply) << std::endl;
    }
    else
    {
        std::vector<std::byte> received;
        auto rres = ch.recv_frame(peer, received, ns);
        if (rres.is_error())
        {
            std::cerr << "[channel-ns] recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-ns] rank 1 received: " << to_string(received) << std::endl;

        std::vector<std::byte> response = to_bytes(to_string(received) + " | ack");
        auto sres = ch.send_frame(peer, response, ns);
        if (sres.is_error())
        {
            std::cerr << "[channel-ns] send error: " << sres.error_message() << std::endl;
            return;
        }
        std::cout << "[channel-ns] rank 1 sent: " << to_string(response) << std::endl;
    }
}

static void
demo_taglease_iprobe(Channel & ch, TagNamespace ns)
{
    int rank = ch.rank();
    int peer = (rank == 0 ? 1 : 0);

    if (rank == 0)
    {
        // 通过租约获取一次性标签，使用显式 tag 发送
        TagLease lease = ns.lease();
        std::vector<std::byte> payload = to_bytes("lease-send: hi");
        auto sres = ch.send_frame(peer, payload, lease.value());
        if (sres.is_error())
        {
            std::cerr << "[tag-lease] send error: " << sres.error_message() << std::endl;
        }
        else
        {
            std::cout << "[tag-lease] rank 0 sent with tag " << lease.value() << ": "
                      << to_string(payload) << std::endl;
        }
        // 主动归还（也可依赖析构自动释放）
        lease.retire();
    }
    else
    {
        // 使用命名空间的非阻塞探测获取确切标签，然后显式接收
        MessageTag probed_tag(0, ns.id());
        bool ready = false;
        while (!ready)
        {
            ready = ns.iprobe(ch.comm(), peer, probed_tag);
        }

        std::vector<std::byte> received;
        auto rres = ch.recv_frame(peer, received, probed_tag.value);
        if (rres.is_error())
        {
            std::cerr << "[tag-lease] recv error: " << rres.error_message() << std::endl;
            return;
        }
        std::cout << "[tag-lease] rank 1 received with tag " << probed_tag.value << ": "
                  << to_string(received) << std::endl;
    }
}

static void
demo_print_tag_stats(Context & ctx, TagNamespace ns)
{
    auto & alloc = ctx.tag_allocator();
    TagNamespace::Stats ns_stats = ns.get_stats();
    auto gstats = alloc.get_global_stats();

    if (ctx.rank() == 0)
    {
        std::cout << "[tag-stats] ns_allocated=" << ns_stats.allocated_count
                  << ", ns_available=" << ns_stats.available_count
                  << ", ns_total_capacity=" << ns_stats.total_capacity
                  << ", ns_range_count=" << ns_stats.range_count
                  << " | global_namespaces=" << gstats.total_namespaces
                  << ", global_allocated_tags=" << gstats.total_allocated_tags
                  << ", global_available_tags=" << gstats.total_available_tags
                  << ", tag_ub=" << gstats.tag_ub_value << std::endl;
    }
}

static void
demo_disallow_any_tag(Channel & ch, TagNamespace ns)
{
    ChannelOptions opts = ch.options();
    opts.allow_any_tag = false;
    ch.set_options(opts);

    // 尝试命名空间接收，应立刻返回错误
    std::vector<std::byte> dummy;
    auto rres = ch.recv_frame((ch.rank() == 0 ? 1 : 0), dummy, ns);
    if (rres.is_error())
    {
        std::cout << "[channel-options] ANY_TAG disabled error: " << rres.error_message()
                  << std::endl;
    }

    // 恢复为允许 ANY_TAG
    opts.allow_any_tag = true;
    ch.set_options(opts);
}

static void
demo_failure_policy_override(ChannelManager & mgr)
{
    ChannelOptions opts;
    opts.use_dedicated_comm = true;
    opts.override_failure_policy = true;
    opts.failure_policy = FailurePolicy::LogAndContinue;
    // 演示失败策略覆盖时的安全错误路径：禁止 ANY_TAG
    opts.allow_any_tag = false;

    auto cres = mgr.create_channel(ChannelId("ch-failure-policy"), opts);
    if (cres.is_error())
    {
        std::cerr << "[manager] create channel error: " << cres.error_message() << std::endl;
        return;
    }
    auto ch = cres.value();
    (void)ch->activate();

    // 使用命名空间接收以触发 ANY_TAG 禁止的错误（不涉及 MPI_FATAL）
    auto & alloc = ch->context().tag_allocator();
    TagNamespace ns = alloc.create_namespace("fp_override_ns");
    int peer = (ch->rank() == 0 ? 1 : 0);
    std::vector<std::byte> out;
    auto rres = ch->recv_frame(peer, out, ns);
    if (rres.is_error())
    {
        std::cout << "[failure-policy] recv error captured (LogAndContinue): "
                  << rres.error_message() << std::endl;
    }

    (void)ch->close();
    (void)mgr.destroy_channel(ChannelId("ch-failure-policy"));
}

static void
demo_manager_and_group(Context & ctx)
{
    auto & mgr = ChannelManager::for_comm(ctx.comm());

    ChannelOptions opts1;
    opts1.use_dedicated_comm = true;
    opts1.default_namespace_name = "channel_default";
    opts1.allow_any_tag = true;

    auto ch_res = mgr.create_channel(ChannelId("ch-main"), opts1);
    if (ch_res.is_error())
    {
        std::cerr << "[manager] create channel error: " << ch_res.error_message() << std::endl;
        return;
    }
    auto ch = ch_res.value();
    (void)ch->activate();

    // 创建命名空间（确保与选项保持一致）
    auto & alloc = ctx.tag_allocator();
    TagNamespace ns = alloc.create_namespace("channel_ns");

    // 展示显式 tag 与命名空间 + ANY_TAG 过滤收发
    demo_channel_explicit_tag(*ch);
    demo_channel_namespace_any(*ch, ns);

    // 展示 TagLease 与 iprobe 的组合：发送端持有租约显式标签，接收端探测标签
    demo_taglease_iprobe(*ch, ns);

    // 展示关闭 ANY_TAG 时的错误路径
    demo_disallow_any_tag(*ch, ns);

    // 统计信息
    const auto & st = ch->stats();
    if (ch->rank() == 0)
    {
        std::cout << "[stats] messages_sent=" << st.messages_sent
                  << ", messages_received=" << st.messages_received
                  << ", bytes_sent=" << st.bytes_sent << ", bytes_received=" << st.bytes_received
                  << std::endl;
    }

    // 打印命名空间与全局标签统计信息
    demo_print_tag_stats(ctx, ns);

    // 管理器与组操作
    ChannelOptions opts2;
    opts2.use_dedicated_comm = true;
    auto ch2_res = mgr.create_channel(ChannelId("ch-aux"), opts2);
    if (!ch2_res.is_error())
    {
        auto ch2 = ch2_res.value();
        ChannelGroup group("channel-group-demo");
        group.add_channel(ch);
        group.add_channel(ch2);
        (void)group.activate_all();
        std::cout << "[group] size=" << group.size()
                  << ", active_count=" << mgr.active_channel_count() << std::endl;
        (void)group.close_all();
    }

    // 覆盖失败策略示例
    demo_failure_policy_override(mgr);

    // 清理
    (void)ch->close();
    (void)mgr.destroy_channel(ChannelId("ch-main"));
    (void)mgr.destroy_channel(ChannelId("ch-aux"));
}

int
main()
{
    Environment env;
    env.set_failure_policy(FailurePolicy::LogAndContinue);
    Context & ctx = Context::for_comm(MPI_COMM_WORLD, env);

    if (ctx.size() < 2)
    {
        if (ctx.rank() == 0)
            std::cerr << "This example requires at least 2 MPI processes." << std::endl;
        return 0;
    }

    demo_manager_and_group(ctx);
    return 0;
}